import tensorflow as tf
import numpy as np

from tensorflow.python.ops import math_ops,array_ops,variable_scope
from tensorflow.python.ops.metrics_impl import mean,_streaming_confusion_matrix

label = np.array([
    [1],
    [1],
    [1],
    [0],
    [0]
])

prediction = np.array([
    [1],
    [1],
    [1],
    [0],
    [1]
])

num_class = 2


def compute_mean_iou(label_, prediction_ ,num_class_, name='ious'):
    with variable_scope.variable_scope(name, 'ious',
                                       (prediction_, label_, None)):
        total_cm, update_op = _streaming_confusion_matrix(label_, prediction_,
                                                          num_class_, None)
        # total_cm = tf.Print(total_cm, ['de total_cm:', total_cm])
        # update_op = tf.Print(update_op, ['de update_op:', update_op])
        """Compute the mean intersection-over-union via the confusion matrix."""
        sum_over_row = math_ops.to_float(math_ops.reduce_sum(total_cm, 0))
        sum_over_col = math_ops.to_float(math_ops.reduce_sum(total_cm, 1))
        cm_diag = math_ops.to_float(array_ops.diag_part(total_cm))
        denominator = sum_over_row + sum_over_col - cm_diag

        # The mean is only computed over classes that appear in the
        # label or prediction tensor. If the denominator is 0, we need to
        # ignore the class.
        num_valid_entries = math_ops.reduce_sum(
            math_ops.cast(
                math_ops.not_equal(denominator, 0), dtype=tf.float32))

        # If the value of the denominator is 0, set it to 1 to avoid
        # zero division.
        denominator = array_ops.where(
            math_ops.greater(denominator, 0), denominator,
            array_ops.ones_like(denominator))
        iou = math_ops.div(cm_diag, denominator)

        return iou , update_op


def case1():
    label_placeholder = tf.placeholder(tf.int64, shape=[5, 1])
    prediction_placeholder = tf.placeholder(tf.int64, shape=[5, 1])

    acc_op, acc_op_up = tf.metrics.accuracy(label_placeholder, prediction_placeholder)

    acc_perclass_op, acc_perclass_up_op = tf.metrics.mean_per_class_accuracy(label_placeholder, prediction_placeholder,
                                                                             num_classes=num_class)

    iou_op, iou_up_op = tf.metrics.mean_iou(label_placeholder, prediction_placeholder, num_classes=num_class)

    ious_op  ,ious_up_op = compute_mean_iou(label_placeholder, prediction_placeholder , num_class)

    with tf.Session() as sess:
        sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
        for i in range(2):
            print(i)
            acc, acc_up, iou, iou_up, acc_perclass, acc_perclass_up , ious ,ious_up= sess.run(
                [acc_op, acc_op_up, iou_op, iou_up_op, acc_perclass_op, acc_perclass_up_op, ious_op , ious_up_op],
                feed_dict={label_placeholder: label, prediction_placeholder: prediction})
            # print(iou_up)
            # print(ious)
            # print(ious_up)

        print("acc: {:.4f} iou: {:.4f} acc per class: {:.4f}  ".format(acc, iou, acc_perclass))
        print("acc per class: {}".format(acc_perclass_up))
        print("ious per class:{}".format(ious))

if __name__ == '__main__':
    case1()